import pandas as pd
import pymysql
import numpy as np
from pathlib import Path

# 数据库配置
db_host = '127.0.0.1'
db_port = 3306
db_user = 'root'
db_password = 'root'
db_database = 'tfj'
db_charset = 'utf8mb4'

csv_path = 'data/goods_images_filtered_10w.csv'
# 新表名配置
table_name = 'goods_images_eval_o2'


def create_table_if_not_exists(cursor, table_name, df):
    """创建表（如果不存在）"""
    # 获取列名和数据类型
    columns = []
    for col in df.columns:
        # 根据数据类型选择合适的MySQL数据类型
        dtype = df[col].dtype
        if dtype == np.int64:
            col_type = 'BIGINT'
        elif dtype == np.float64:
            col_type = 'DOUBLE'
        else:
            # 对于字符串类型，使用TEXT以确保能存储长文本
            col_type = 'TEXT'
        columns.append(f"`{col}` {col_type}")

    # 创建表的SQL语句
    create_table_sql = f"""
    CREATE TABLE IF NOT EXISTS `{table_name}` (
        {', '.join(columns)}
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
    """
    cursor.execute(create_table_sql)


def batch_insert_data(cursor, table_name, df, batch_size=1000):
    """批量插入数据"""
    # 构建INSERT语句
    columns = [f"`{col}`" for col in df.columns]
    placeholders = ', '.join(['%s'] * len(df.columns))
    insert_sql = f"INSERT INTO `{table_name}` ({', '.join(columns)}) VALUES ({placeholders})"
    
    # 准备数据
    values = df.replace({np.nan: None}).values.tolist()
    
    # 批量插入
    total_rows = len(values)
    for i in range(0, total_rows, batch_size):
        batch = values[i:i + batch_size]
        cursor.executemany(insert_sql, batch)
        print(f"  -> 已插入 {min(i + batch_size, total_rows)}/{total_rows} 行")


def upload_csv_to_database(csv_path, table_name, db_params):
    """
    主函数：读取CSV文件并上传到数据库
    """
    # 步骤 1: 读取CSV文件
    try:
        print(f"步骤 1: 正在从 '{csv_path}' 读取CSV文件...")
        df = pd.read_csv(csv_path, encoding='utf-8-sig')
        print(f" -> CSV读取成功，共 {len(df)} 行。")
        print(f" -> 列名: {list(df.columns)}")

    except Exception as e:
        print(f"读取CSV文件时出错: {e}")
        return

    # 步骤 2: 连接数据库并上传数据
    print(f"\n步骤 2: 正在连接数据库并上传到表 '{table_name}'...")
    try:
        # 创建数据库连接
        conn = pymysql.connect(
            host=db_params['host'],
            port=db_params['port'],
            user=db_params['user'],
            password=db_params['password'],
            database=db_params['database'],
            charset=db_params['charset']
        )
        
        with conn.cursor() as cursor:
            if not df.empty:
                # 删除已存在的表
                cursor.execute(f"DROP TABLE IF EXISTS `{table_name}`")
                # 创建新表
                create_table_if_not_exists(cursor, table_name, df)
                # 插入数据
                batch_insert_data(cursor, table_name, df)
                print(f" -> '{table_name}' 上传成功。")
            else:
                print(f" -> CSV文件为空，跳过上传。")

            # 提交事务
            conn.commit()

    except Exception as e:
        import traceback
        print(f"\n数据库操作失败: {e}")
        traceback.print_exc()
        if 'conn' in locals():
            conn.rollback()  # 发生错误时回滚
        return
    finally:
        if 'conn' in locals():
            conn.close()

    print("\n🎉 CSV文件上传完成！")


# --- 程序入口 ---
if __name__ == "__main__":
    # 数据库参数
    db_params = {
        'host': db_host,
        'port': db_port,
        'user': db_user,
        'password': db_password,
        'database': db_database,
        'charset': db_charset
    }
    
    # 调用主函数
    upload_csv_to_database(csv_path, table_name, db_params)







